/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/

/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define SPYDLL_BUILD_DLL
#define ENABLE_AUTOEJECT		1
#define TRACE_LEVEL	2	// warning level


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "SpyDllInt.h"

// standard headers
#include <windows.h>
#include <tchar.h>
#include <psapi.h>
#include <tlhelp32.h>
#include <stdio.h> //<-- TO BE REMOVED

// project's headers
#include "Assert.h"
#include "Condition.h"
#include "Hooks.h"
#include "SpyDll.h"
#include "Trace.h"


/******************************************************************/
/* Internal constants                                             */
/******************************************************************/

#define SOFT_AUTOEJECT_THRESHOLD	10
#define HARD_AUTOEJECT_THRESHOLD	50

#define MAX_SERVER_DELAY		10*60*1000 /*ms*/

#define STACK_UNROLL_LENGTH		32



/******************************************************************/
/* Internal data                                                  */
/******************************************************************/

extern DWORD   	sg_dwServerProcessID;
extern PROC	sg_pfnNotifyProc;
extern HINSTANCE g_hInstance ;
extern MODULEINFO g_miInstance ;


/******************************************************************/
/* Exported function : CallRemoteFuncByHandle                     */
/******************************************************************/

BOOL SpyInt_CallRemoteFuncByHandle (HANDLE hRemoteProcess, PROC pfnRemoteFunc, 
				    void * pData, int nDataSize, DWORD * pdwExitCode)
{
  HANDLE	hRemoteThread ;
  void		*pRemoteData ;
  BOOL		bResult = FALSE ;

  //  ASSERT (dwRemoteProcessID!=0) ;
  ASSERT (pfnRemoteFunc!=NULL) ;
  ASSERT (pdwExitCode!=NULL) ;
  ASSERT (hRemoteProcess!=NULL) ;
  
  if( nDataSize>0 )
    {
      pRemoteData = VirtualAllocEx (hRemoteProcess,  NULL, nDataSize, 
				    MEM_COMMIT, PAGE_READWRITE) ;
      
      if( ! pRemoteData ) {	
	CloseHandle (hRemoteProcess) ;
	return FALSE ;      
      }
      
      if( ! WriteProcessMemory (hRemoteProcess, pRemoteData,
				pData, nDataSize, NULL) )
	{
	  VirtualFreeEx (hRemoteProcess, pRemoteData, 0, MEM_RELEASE) ;
	  CloseHandle (hRemoteProcess) ;
	  return FALSE ;     
	}	  
    }
  else
    {
      pRemoteData = pData ;
    }
  
  TRACE_INFO (TEXT("Process = 0x%08X\n"), hRemoteProcess) ;

  hRemoteThread = CreateRemoteThread (hRemoteProcess, NULL, 0, 
				      (LPTHREAD_START_ROUTINE)pfnRemoteFunc, 
				      pRemoteData, 0, NULL) ;
  
  if( hRemoteThread )
    {
      if( WAIT_OBJECT_0==WaitForSingleObject(hRemoteThread, MAX_SERVER_DELAY) )
	{    
	  if( GetExitCodeThread (hRemoteThread, pdwExitCode) )
	    bResult = ( *pdwExitCode != STILL_ACTIVE ) ;
	}
      else 
	{
	  TRACE_ERROR (TEXT("Remote thread seems freezed, I kill him.\n")) ;
	  TerminateThread (hRemoteThread, 0) ;
	}
      
      CloseHandle (hRemoteThread) ;
    }
  else
    {
      TRACE_ERROR (TEXT("CreateRemoteThread failed (error=%u)\n"), GetLastError()) ;
    }
  
  if( nDataSize>0 ) VirtualFreeEx (hRemoteProcess, pRemoteData, 0, MEM_RELEASE) ;

  TRACE_INFO (TEXT("returned %u\n"), *pdwExitCode) ;
          
  return bResult ;
}


/******************************************************************/
/* Exported function : CallRemoteFuncByPID                        */
/******************************************************************/

BOOL SpyInt_CallRemoteFuncByPID (DWORD dwRemoteProcessID, PROC pfnRemoteFunc, 
				 void * pData, int nDataSize, 
				 DWORD * pdwExitCode)
{
  HANDLE	hRemoteProcess ;
  BOOL		bResult = FALSE ;
  DWORD		dwDesiredAccess = PROCESS_CREATE_THREAD	| PROCESS_VM_OPERATION | PROCESS_VM_WRITE ;

  //  ASSERT (dwRemoteProcessID!=0) ;
  ASSERT (pfnRemoteFunc!=NULL) ;
  ASSERT (pdwExitCode!=NULL) ;

  TRACE_INFO (TEXT("Process = %u\n"), dwRemoteProcessID) ;

  // is it the current process ?
  if( dwRemoteProcessID==GetCurrentProcessId() )
    {
      *pdwExitCode = pfnRemoteFunc (pData) ;
      bResult = TRUE ;
    }
  // is it a remote process ?
  else
    {
      hRemoteProcess = OpenProcess (dwDesiredAccess, FALSE, dwRemoteProcessID) ;
      
      if( hRemoteProcess )
	{	    	
	  bResult = SpyInt_CallRemoteFuncByHandle (hRemoteProcess, pfnRemoteFunc, 
						   pData, nDataSize, pdwExitCode);
	  
	  CloseHandle (hRemoteProcess) ;
	}
      else
	{
	  if( GetLastError()==ERROR_ACCESS_DENIED )
	    {
	      TRACE_INFO (TEXT("Insufficient rights to OpenProcess %u.\n"), dwRemoteProcessID) ;	  
	    }
	  else
	    {
	      TRACE_WARNING (TEXT("OpenProcess failed (pid=%u, error=%u)\n"), 
			     dwRemoteProcessID, GetLastError()) ;
	    }
	}
    }
      
  return bResult ;
}


/******************************************************************/
/* Exported function : AsciiToUnicode                             */
/******************************************************************/

TCHAR * SpyInt_AsciiToUnicode (LPCSTR szSrc) 
{
  int nDstLen ;
  TCHAR * szDst ;

  if( ! szSrc ) {
    TRACE_WARNING (TEXT("szSrc==NULL\n")) ;
    return NULL ;
  }
  
  // calc needed size
  nDstLen =  MultiByteToWideChar (CP_THREAD_ACP, 0, szSrc, -1, NULL, 0) ;

  if( ! nDstLen > 200 )
    TRACE_WARNING (TEXT("Big string (length=%u)\n"), nDstLen) ;
 
  // alloc
  szDst = malloc (nDstLen*sizeof(TCHAR)) ;

  if( ! szSrc ) {
    TRACE_ERROR (TEXT("malloc failed\n")) ;
    return NULL ;
  }

  // convert
  MultiByteToWideChar (CP_THREAD_ACP, 0, szSrc, -1, szDst, nDstLen) ;

  return szDst ;
}


/******************************************************************/
/* Exported function : CatchCall                                  */
/******************************************************************/

DWORD SpyInt_CatchCall (UINT nReason, LPCTSTR szFormat, ...)
{
  SDNCATCHCALL *pMessage ;
  UINT		nMessageSize ;
  DWORD		dwResponse ;
  va_list	va ;

  // if server isn't totally initialized
  // then accept call
  if( sg_pfnNotifyProc==NULL || sg_dwServerProcessID==0 ) {
    TRACE_INFO (TEXT("Lost a call : %s\n"), Reason_GetName(nReason)) ;
    return RULE_ACCEPT ;
  }

  va_start (va, szFormat) ;

  nMessageSize = sizeof(SDNCATCHCALL) + Condition_CalcDataSizeV(szFormat,va) ;
 
  pMessage		= malloc (nMessageSize) ;
  pMessage->hdr.dwCode	= SDN_CATCH_CALL ;
  pMessage->hdr.dwProcessId	= GetCurrentProcessId() ;
  pMessage->cond.nSize	= nMessageSize - (sizeof(SDNCATCHCALL)-sizeof(CONDITION)) ;
    
  Condition_SetV (&pMessage->cond, nReason, szFormat, va) ;

  dwResponse = SpyInt_NotifyServer ((SDNMHDR*)pMessage, nMessageSize) ;

  free (pMessage) ;

  va_end (va) ;

  if( dwResponse==RULE_KILLPROCESS )
    ExitProcess (0) ;

#if ENABLE_AUTOEJECT  
  // auto eject
  if( dwResponse==SPYINT_ERROR ) 
    {
      static int nCount ;

      nCount++ ;

      if( nCount==SOFT_AUTOEJECT_THRESHOLD )
	{
	  TRACE_ERROR (TEXT("Too many errors : auto-eject (soft)\n")) ;
	  //SpyDll_EjectFromProcess (GetCurrentProcessId()) ;
	  if( QueueUserWorkItem ((LPTHREAD_START_ROUTINE)FreeLibrary, g_hInstance, 
				 WT_EXECUTEDEFAULT|WT_EXECUTELONGFUNCTION) )
	    TRACE_INFO (TEXT("Auto-eject queued\n")) ;
	  else
	    TRACE_ERROR (TEXT("Failed to queue auto-eject (error=%d)\n"), GetLastError()) ;
	}

      if( nCount==HARD_AUTOEJECT_THRESHOLD )
	{
	  TRACE_ERROR (TEXT("Too many errors : auto-eject (hard)\n")) ;
	  FreeLibraryAndExitThread (g_hInstance,0) ;
	}
    }
#endif

  if( dwResponse==SPYINT_ERROR ) 
    {
      TRACE_ERROR (TEXT("Failed to call filter function (reason=%s)\n"), 
		   Reason_GetName(nReason)) ;
      dwResponse = RULE_ACCEPT ;
    }

  return dwResponse ;
}


/******************************************************************/
/* Exported function : NotifyServer                               */
/******************************************************************/

DWORD SpyInt_NotifyServer (SDNMHDR * pMessage, DWORD dwSize) 
{
  DWORD		dwExitCode ;

  //  DbgPrint (TEXT("---- Code = %d\n"), pMessage->dwCode) ;

  if( ! sg_pfnNotifyProc || !sg_dwServerProcessID ) 
    {
      TRACE_WARNING (TEXT("DLL is not correctly initialized\n")) ;
      return SPYINT_ERROR ;
    }
  
  if( !SpyInt_CallRemoteFuncByPID (sg_dwServerProcessID, (PROC)sg_pfnNotifyProc, 
				   pMessage, dwSize, &dwExitCode) )
    {
      TRACE_WARNING (TEXT("Failed to notify server\n")) ;

      if( ERROR_INVALID_PARAMETER==GetLastError() )
	{
	  TRACE_WARNING (TEXT("Server seems crashed, resetting server params\n")) ;
	  sg_dwServerProcessID = 0 ;
	  sg_pfnNotifyProc = NULL ;
	}
      
      dwExitCode = SPYINT_ERROR ;
    }
 
  return dwExitCode ;
}


/******************************************************************/
/* Exported function : SuspendThread                              */
/******************************************************************/

BOOL	SpyInt_SuspendThread (DWORD dwThreadID,
			      MEMAREA * pForbidden, UINT nForbidden)
{
  BOOL	bOK = FALSE ;
  int	iSuspendCount ;
  UINT	iAttempt, iArea ;
  DWORD nBeginDll, nEndDll ;
  HANDLE hThread = OpenThread (THREAD_SUSPEND_RESUME|THREAD_GET_CONTEXT, 
			       FALSE, dwThreadID) ;  
  
  ASSERT (dwThreadID!=GetCurrentThreadId()) ;
  
  iSuspendCount = SuspendThread (hThread) ;

  if( iSuspendCount<0 ) {
    TRACE_ERROR (TEXT("Failed to suspend thread %d\n"), dwThreadID) ;
    CloseHandle(hThread) ;
    return FALSE ;     
  }

  if( iSuspendCount>0 ) {
    TRACE_WARNING (TEXT("Thread is already suspended (count=%d)\n"), iSuspendCount) ;
    return TRUE ;         
  }

  nBeginDll = (DWORD) g_miInstance.lpBaseOfDll ;
  nEndDll = nBeginDll + g_miInstance.SizeOfImage ;
  
  for( iAttempt=0 ; iAttempt<10 && !bOK ; iAttempt++ )
    {
      CONTEXT	context = {0} ; 
      DWORD	aAddresses[STACK_UNROLL_LENGTH] ;
      DWORD	iAddress ;
      
      context.ContextFlags = CONTEXT_CONTROL ;
      if( ! GetThreadContext (hThread, &context) ) {
	TRACE_WARNING (TEXT("GetThreadContext failed (error=%u)\n"),
		       GetLastError()) ;
	CloseHandle(hThread) ;
	return TRUE ;
      }	  

      aAddresses[0] = context.Eip ;
      for( iAddress=1 ; iAddress<STACK_UNROLL_LENGTH ; iAddress++)
	{
	  DWORD * pStackLoc = (DWORD*)context.Esp + (iAddress-1) ;
	  if( IsBadReadPtr(pStackLoc,sizeof(DWORD)) )
	    aAddresses[iAddress] = 0 ;
	  else
	    aAddresses[iAddress] = *pStackLoc ;
	}
      
      TRACE_INFO (TEXT("Thread %u, EIP=0x%08X\n"), dwThreadID, aAddresses[0]) ;
      TRACE_INFO (TEXT("Thread %u, Stack=0x%08X\n"), dwThreadID, aAddresses[1]) ;
	  
      bOK = TRUE ;
      
      for( iAddress=0 ; iAddress<STACK_UNROLL_LENGTH && bOK ; iAddress++ )
	{
	  for( iArea=0 ; iArea<nForbidden ; iArea++ )
	    {
	      DWORD nBeginArea = (DWORD)pForbidden[iArea].pBase ;
	      DWORD nEndArea = nBeginArea + pForbidden[iArea].nSize ;
	      
	      if( aAddresses[iAddress]>=nBeginDll && aAddresses[iAddress]<nEndDll  )
		{
		  TRACE_WARNING (TEXT("Thread %u, unroll %d -> 0x%08X\n"), 
				 dwThreadID, iAddress, aAddresses[iAddress]) ;
		  TRACE_WARNING (TEXT("Thread %u in current DLL (0x%08X-0x%08X)\n"), 
				 dwThreadID, nBeginDll, nEndDll) ;
		  bOK = FALSE ;
		  break ;
		}
	  
	      if( aAddresses[iAddress]>=nBeginArea && aAddresses[iAddress]<nEndArea )
		{
		  TRACE_WARNING (TEXT("Thread %u, unroll %d -> 0x%08X\n"), 
				 dwThreadID, iAddress, aAddresses[iAddress]) ;
		  TRACE_WARNING (TEXT("Thread %u in forbidden area %d (0x%08X-0x%08X)\n"), 
				 dwThreadID, iArea, nBeginArea, nEndArea) ;
		  bOK = FALSE ;
		  break ;
		}
	    }
	}
      
      // Not OK => let some time for EIP to move
      if( !bOK )
	{
	  iSuspendCount = ResumeThread(hThread) ;
	  if( 1!=iSuspendCount ) {
	    TRACE_ERROR (TEXT("Failed to resume thread (count=%d)\n"), iSuspendCount) ;
	    break ;
	  }

	  Sleep (100) ;

	  iSuspendCount = SuspendThread(hThread) ;
	  if( 0!=iSuspendCount ) {
	    TRACE_ERROR (TEXT("Failed to suspend thread (count=%d)\n"), iSuspendCount) ;
	    break ;
	  }
	}
    }
  
  if( !bOK )
    TRACE_ERROR (TEXT("Thread %u is in a forbidden area or state\n"), dwThreadID) ;
  
  CloseHandle(hThread) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function :                                            */
/******************************************************************/

BOOL SpyInt_SuspendOtherThreads (MEMAREA * pForbidden, UINT nForbidden)
{
  HANDLE hSnapShot ; 
  THREADENTRY32 te32 ;
  
  hSnapShot = CreateToolhelp32Snapshot (TH32CS_SNAPTHREAD, 0) ;
  if( hSnapShot==INVALID_HANDLE_VALUE ) return FALSE ;

  te32.dwSize = sizeof(te32) ;

  if( ! Thread32First (hSnapShot, &te32) )
    {
      CloseHandle (hSnapShot) ;
      return FALSE ;
    }
  
  do {
    
    if( te32.th32OwnerProcessID == GetCurrentProcessId() ) 
      if( te32.th32ThreadID != GetCurrentThreadId() )
	SpyInt_SuspendThread (te32.th32ThreadID, pForbidden, nForbidden) ;
 
  }
  while( Thread32Next(hSnapShot,&te32) ) ;
  
  CloseHandle (hSnapShot) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function : ResumeThread                               */
/******************************************************************/

BOOL	SpyInt_ResumeThread (DWORD dwThreadID)
{
  HANDLE hThread = OpenThread (THREAD_SUSPEND_RESUME|THREAD_GET_CONTEXT, 
			       FALSE, dwThreadID) ;  

  ASSERT (dwThreadID!=GetCurrentThreadId()) ;

  ResumeThread (hThread) ;
    
  CloseHandle(hThread) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function :                                            */
/******************************************************************/

BOOL SpyInt_ResumeOtherThreads (MEMAREA * pForbidden, UINT nForbidden)
{
  HANDLE hSnapShot ; 
  THREADENTRY32 te32 ;
  
  hSnapShot = CreateToolhelp32Snapshot (TH32CS_SNAPTHREAD, 0) ;
  if( hSnapShot==INVALID_HANDLE_VALUE ) return FALSE ;

  te32.dwSize = sizeof(te32) ;

  if( ! Thread32First (hSnapShot, &te32) )
    {
      CloseHandle (hSnapShot) ;
      return FALSE ;
    }
  
  do {
    
    if( te32.th32OwnerProcessID == GetCurrentProcessId() ) 
      if( te32.th32ThreadID != GetCurrentThreadId() )
	SpyInt_ResumeThread (te32.th32ThreadID) ;
 
  }
  while( Thread32Next(hSnapShot,&te32) ) ;
  
  CloseHandle (hSnapShot) ;

  return TRUE ;
}

